package com.quyu.service;

import com.fasterxml.jackson.core.type.TypeReference;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

public class WebSocketServer extends AbstractWebSocketHandler {
    private final Set<WebSocketSession> sessions = Collections.synchronizedSet(new HashSet<>());
    private final AtomicInteger userCount = new AtomicInteger(0);
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.add(session);
        userCount.incrementAndGet();
        broadcastUserCount();
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        try {
            // 1. 记录收到的原始消息
            String payload = message.getPayload();
            System.out.println("Received message from " + session.getId() + ": " + payload);

            // 2. 解析JSON消息
            ObjectMapper objectMapper = new ObjectMapper();
            Map<String, Object> messageMap;

            try {
                messageMap = objectMapper.readValue(payload, new TypeReference<Map<String, Object>>() {});
            } catch (IOException e) {
                System.err.println("Failed to parse JSON message: " + e.getMessage());
                session.sendMessage(new TextMessage("{\"error\":\"Invalid JSON format\"}"));
                return;
            }

            // 3. 验证必要字段
            if (!messageMap.containsKey("type") || !messageMap.containsKey("sender")) {
                System.err.println("Missing required fields in message");
                session.sendMessage(new TextMessage("{\"error\":\"Missing required fields (type/sender)\"}"));
                return;
            }

            String type = (String) messageMap.get("type");
            String sender = (String) messageMap.get("sender");
            String content = messageMap.containsKey("content") ? (String) messageMap.get("content") : "";

            // 4. 根据消息类型处理
            switch (type) {
                case "JOIN":
                    // 新用户加入
                    System.out.println(sender + " joined the chat");
                    broadcastSystemMessage(sender + " 加入了聊天室");
                    break;

                case "LEAVE":
                    // 用户离开
                    System.out.println(sender + " left the chat");
                    broadcastSystemMessage(sender + " 离开了聊天室");
                    break;

                case "MESSAGE":
                    // 普通聊天消息
                    if (content == null || content.trim().isEmpty()) {
                        System.err.println("Empty message content");
                        return;
                    }
                    System.out.println(sender + " says: " + content);

                    // 构建返回消息
                    Map<String, Object> response = new HashMap<>();
                    response.put("type", "MESSAGE");
                    response.put("sender", sender);
                    response.put("content", content);
                    response.put("timestamp", System.currentTimeMillis());

                    // 广播消息给所有客户端
                    broadcastMessage(objectMapper.writeValueAsString(response));
                    break;

                case "TYPING":
                    // 用户正在输入
                    Map<String, Object> typingResponse = new HashMap<>();
                    typingResponse.put("type", "TYPING");
                    typingResponse.put("sender", sender);
                    broadcastMessage(objectMapper.writeValueAsString(typingResponse));
                    break;

                default:
                    System.err.println("Unknown message type: " + type);
                    session.sendMessage(new TextMessage("{\"error\":\"Unknown message type\"}"));
            }

        } catch (Exception e) {
            System.err.println("Error handling message: " + e.getMessage());
            e.printStackTrace();
            try {
                session.sendMessage(new TextMessage("{\"error\":\"Server error: " + e.getMessage() + "\"}"));
            } catch (IOException ioException) {
                System.err.println("Failed to send error message: " + ioException.getMessage());
            }
        }
    }

    // 广播系统消息
    private void broadcastSystemMessage(String message) throws IOException {
        Map<String, Object> systemMessage = new HashMap<>();
        systemMessage.put("type", "SYSTEM");
        systemMessage.put("content", message);
        systemMessage.put("timestamp", System.currentTimeMillis());

        broadcastMessage(new ObjectMapper().writeValueAsString(systemMessage));
    }

    // 广播普通消息
    private void broadcastMessage(String message) throws IOException {
        synchronized (sessions) {
            for (WebSocketSession s : sessions) {
                try {
                    if (s.isOpen()) {
                        s.sendMessage(new TextMessage(message));
                    }
                } catch (IOException e) {
                    System.err.println("Failed to send message to " + s.getId() + ": " + e.getMessage());
                    sessions.remove(s);
                    try {
                        s.close();
                    } catch (IOException ex) {
                        System.err.println("Failed to close session: " + ex.getMessage());
                    }
                }
            }
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        sessions.remove(session);
        userCount.decrementAndGet();
        broadcastUserCount();
    }

    private void broadcastUserCount() throws IOException {
        broadcast(objectMapper.writeValueAsString(Map.of(
                "type", "USER_COUNT",
                "sender", "系统",
                "content", userCount.toString()
        )));
    }

    private void broadcast(String message) throws IOException {
        synchronized (sessions) {
            for (WebSocketSession session : sessions) {
                if (session.isOpen()) {
                    try {
                        session.sendMessage(new TextMessage(message));
                    } catch (IOException e) {
                        sessions.remove(session);
                    }
                }
            }
        }
    }
}